import numpy as np
import math
import gym
from gym import spaces
from geopy.distance import geodesic
from getDCSdata import get_dcs_data
from keyboard_control import (
    pitch_up, pitch_down, roll_left, roll_right, throttle_up, throttle_down,
    flapsland_control, gearson_control, airbrake_control, foot_left, foot_right,
    activate_dcs_window
)
import logging
import time
import matplotlib.pyplot as plt
import csv
import os

# 日志配置
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("results/descent_env.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 跑道参数（科布列季机场，Su-25T）
RUNWAY_LAT = 41.929631
RUNWAY_LONG = 41.86325
RUNWAY_ALT = 18.0
APPROACH_SPEED = 79.17  # 285 km/h ≈ 79.17 m/s
LANDING_SPEED = 65.28  # 235 km/h ≈ 65.28 m/s
RUNWAY_HEADING = 76
TRUE_HEADING = 64.0
GLIDE_ANGLE = math.radians(3.0)  # 3°下滑角
INTERCEPT_DISTANCE = 15000  # 15km
AOA_THRESHOLD = 22.0
SPEED_UPPER_LIMIT = 83.33  # 300 km/h ≈ 83.33 m/s
SPEED_LOWER_LIMIT = 55.56  # 200 km/h ≈ 55.56 m/s
LOW_ALT_THRESHOLD = 350.0  # 起落架和襟翼触发高度
CRASH_ALT_THRESHOLD = RUNWAY_ALT - 50.0  # 坠毁高度 -50m

# CSV文件路径
GLIDE_PATH_CSV = "results/glide_path.csv"
GLIDE_DEVIATION_CSV = "results/glide_deviation.csv"
FLIGHT_DATA_CSV = "results/flight_data.csv"
TRAINING_LOG_CSV = "results/descent_training_log.csv"


class DescentEnv(gym.Env):
    def __init__(self):
        super().__init__()
        # 动作空间：6维连续动作（俯仰、滚转、油门、襟翼、减速板、脚舵）
        self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(6,), dtype=np.float32)
        # 观察空间：17维状态（经纬高、速度、姿态、偏差等）
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(17,), dtype=np.float32)

        # 初始化变量
        self.steps = 0
        self.current_episode = 0
        self.total_reward = 0
        self.is_on_ground = False
        self.last_distance = None
        self.last_altitude = None
        self.runway_point = (RUNWAY_LAT, RUNWAY_LONG)
        self.glide_path = []
        self.target_point = None
        self.smoothed_action = [0] * 6
        self.last_action = [0] * 6
        self.flight_trajectory = []
        self.flight_data = []

        # 初始化CSV文件
        self._init_csv_files()
        # 生成下滑道
        self._generate_glide_path()

    def _init_csv_files(self):
        """初始化CSV文件并写入表头"""
        # 训练日志
        with open(TRAINING_LOG_CSV, mode='w', newline='') as file:
            writer = csv.writer(file)
            writer.writerow([
                "Episode", "Step", "Reward", "Total Reward", "Alt", "IAS",
                "Heading Error", "AoA", "Vertical Speed", "Gears", "Flaps"
            ])

        # 飞行数据
        with open(FLIGHT_DATA_CSV, mode='w', newline='') as file:
            writer = csv.writer(file)
            writer.writerow([
                "Episode", "Step", "Long", "Lat", "Alt", "IAS", "AoA",
                "Pitch", "Bank", "Heading", "Vertical Speed", "Gears", "Flaps",
                "AirBrake", "Fuel", "Dist to Runway"
            ])

        # 下滑道偏差
        with open(GLIDE_DEVIATION_CSV, mode='w', newline='') as file:
            writer = csv.writer(file)
            writer.writerow([
                "Episode", "Step", "Long", "Lat", "Alt", "Path Deviation",
                "Alt Error", "Path Angle Error"
            ])

    def _generate_glide_path(self):
        """生成3°下滑道并保存到CSV"""
        gsip = geodesic(kilometers=INTERCEPT_DISTANCE / 1000).destination(
            self.runway_point, RUNWAY_HEADING + 180
        )
        gsip_alt = RUNWAY_ALT + INTERCEPT_DISTANCE * math.tan(GLIDE_ANGLE)
        self.glide_path = []
        num_points = 150
        distances = np.linspace(0, INTERCEPT_DISTANCE, num_points)
        for dist in distances:
            point = geodesic(kilometers=dist / 1000).destination(
                (gsip.latitude, gsip.longitude), (RUNWAY_HEADING ) % 360
            )
            alt = RUNWAY_ALT + (INTERCEPT_DISTANCE - dist) * math.tan(GLIDE_ANGLE)
            self.glide_path.append({"lat": point.latitude, "lon": point.longitude, "alt": alt})

        # 保存下滑道到CSV
        with open(GLIDE_PATH_CSV, mode='w', newline='') as file:
            writer = csv.writer(file)
            writer.writerow(["Lat", "Lon", "Alt"])
            for p in self.glide_path:
                writer.writerow([p["lat"], p["lon"], p["alt"]])

        logger.info(f"Generated glide path with {num_points} points, saved to {GLIDE_PATH_CSV}")

    def _plot_trajectory(self):
        """绘制飞行轨迹与下滑道对比"""
        if not self.flight_trajectory:
            return

        traj_lats, traj_lons, traj_alts = zip(*self.flight_trajectory)
        path_lats = [p["lat"] for p in self.glide_path]
        path_lons = [p["lon"] for p in self.glide_path]
        path_alts = [p["alt"] for p in self.glide_path]

        # 轨迹图（经纬度）
        plt.figure(figsize=(12, 7))
        plt.plot(path_lons, path_lats, 'b-', label='Glide Path')
        plt.plot(traj_lons, traj_lats, 'g--', label='Flight Trajectory')
        plt.scatter([RUNWAY_LONG], [RUNWAY_LAT], color='red', label='Runway')
        plt.xlabel('Longitude')
        plt.ylabel('Latitude')
        plt.title(f'Flight Trajectory vs Glide Path (Episode {self.current_episode})')
        plt.legend()
        plt.grid(True)
        plt.savefig(f'results/trajectory_episode_{self.current_episode}.png')
        plt.close()

        # 高度剖面图
        plt.figure(figsize=(12, 4))
        plt.plot(range(len(path_alts)), path_alts, 'b-', label='Glide Path Altitude')
        plt.plot(range(len(traj_alts)), traj_alts, 'g--', label='Flight Altitude')
        plt.axhline(RUNWAY_ALT + LOW_ALT_THRESHOLD, color='purple', linestyle='--', label='350m Altitude')
        plt.xlabel('Path Point Index')
        plt.ylabel('Altitude (m)')
        plt.title(f'Altitude Profile (Episode {self.current_episode})')
        plt.legend()
        plt.grid(True)
        plt.savefig(f'results/altitude_profile_episode_{self.current_episode}.png')
        plt.close()

        logger.info(f"Trajectory and altitude profile saved for episode {self.current_episode}")

    def reset(self):
        """重置环境"""
        self.current_episode += 1
        self.steps = 0
        self.total_reward = 0
        self.last_distance = None
        self.last_altitude = None
        self.is_on_ground = False
        self.target_point = None
        self.flight_trajectory = []
        self.flight_data = []
        self.smoothed_action = [0] * 6
        self.last_action = [0] * 6

        try:
            activate_dcs_window()
        except Exception as e:
            logger.error(f"Failed to activate DCS window: {str(e)}")
            data = self._get_fallback_data()
            self.state = self._extract_state(data)
            self.is_on_ground = data["LLA_balt"] - RUNWAY_ALT < 50.0
            self._write_flight_data(data)
            self._write_deviation_data(data)
            return self.state

        data = self._get_dcs_data_with_retry(max_retries=3)
        if data is None:
            logger.warning("Using fallback data after retries")
            data = self._get_fallback_data()

        self.state = self._extract_state(data)
        self.is_on_ground = data["LLA_balt"] - RUNWAY_ALT < 50.0
        self._write_flight_data(data)
        self._write_deviation_data(data)
        logger.info(f"Episode {self.current_episode} started, on_ground={self.is_on_ground}")
        return self.state

    def _get_dcs_data_with_retry(self, max_retries=3):
        """尝试获取DCS数据，带重试机制"""
        for attempt in range(max_retries):
            try:
                data = get_dcs_data()
                if data is not None:
                    return data
                retry_delay = 0.3 + 0.2 * attempt
                logger.warning(
                    f"Data acquisition failed (attempt {attempt + 1}/{max_retries}), retrying after {retry_delay}s")
                time.sleep(retry_delay)
            except Exception as e:
                logger.error(f"Data acquisition error (attempt {attempt + 1}/{max_retries}): {str(e)}")
                time.sleep(retry_delay)
        return None

    def _get_fallback_data(self):
        """返回模拟数据"""
        logger.info("Returning fallback data")
        return {
            "LLA_lat": RUNWAY_LAT + 0.15,
            "LLA_long": RUNWAY_LONG + 0.15,
            "LLA_balt": 800.0,
            "IAS": APPROACH_SPEED,
            "heading_deg": TRUE_HEADING,
            "pitch_deg": 0.0,
            "bank_deg": 0.0,
            "AoA": 8.0,
            "VV": [0, 0, -2.5],
            "Fuel": 0.5,
            "Gears": 0.0,
            "Flaps": 0.0,
            "AirBrake": 0.0,
            "Omega": [0.0, 0.0, 0.0]
        }

    def _extract_state(self, data):
        """提取状态向量"""
        current_pos = (data["LLA_lat"], data["LLA_long"])
        # 找到沿下滑道方向的投影点
        nearest_point, path_deviation = self._find_nearest_glide_point(current_pos)
        self.target_point = nearest_point

        dist_to_runway = geodesic(current_pos, self.runway_point).meters
        alt_error = data["LLA_balt"] - nearest_point["alt"]

        # 计算目标航向（指向最近下滑道点）
        lat1, lon1 = math.radians(data["LLA_lat"]), math.radians(data["LLA_long"])
        lat2, lon2 = math.radians(nearest_point["lat"]), math.radians(nearest_point["lon"])
        dlon = lon2 - lon1
        y = math.sin(dlon) * math.cos(lat2)
        x = math.cos(lat1) * math.sin(lat2) - math.sin(lat1) * math.cos(lat2) * math.cos(dlon)
        target_heading = math.degrees(math.atan2(y, x)) % 360
        heading_error = min(
            abs(target_heading - data["heading_deg"]),
            360 - abs(target_heading - data["heading_deg"])
        )

        # 计算路径角度误差
        target_vec = np.array([nearest_point["lat"] - current_pos[0], nearest_point["lon"] - current_pos[1]])
        heading_vec = np.array(
            [math.cos(math.radians(data["heading_deg"])), math.sin(math.radians(data["heading_deg"]))])
        path_angle_error = math.degrees(math.acos(np.clip(
            np.dot(target_vec, heading_vec) / (np.linalg.norm(target_vec) * np.linalg.norm(heading_vec) + 1e-6), -1, 1
        )))

        state = np.array([
            data["LLA_lat"], data["LLA_long"], data["LLA_balt"],
            data["IAS"], data["pitch_deg"], data["bank_deg"],
            data["AoA"], data["VV"][2], heading_error,
            alt_error, dist_to_runway, data["Fuel"],
            data["Omega"][0], data["Omega"][1], data["Omega"][2],
            path_deviation, path_angle_error
        ], dtype=np.float32)

        assert len(state) == self.observation_space.shape[
            0], f"State length {len(state)} does not match observation space {self.observation_space.shape[0]}"
        return state

    def _find_nearest_glide_point(self, current_pos):
        """找到下滑道上与当前位置的投影点"""
        min_dist = float('inf')
        nearest_point = None
        for p in self.glide_path:
            dist = geodesic(current_pos, (p["lat"], p["lon"])).meters
            if dist < min_dist:
                min_dist = dist
                nearest_point = p
        return nearest_point, min_dist

    def _calculate_reward(self, state, current_distance, data):
        """计算奖励"""
        reward = 0
        path_deviation = state[15]
        path_angle_error = state[16]
        alt_error = state[9]
        heading_error = state[8]
        speed_error = abs(state[3] - APPROACH_SPEED)

        # 动态权重
        w_altitude = 1.2 if current_distance < 3000 else 0.8
        w_deviation = 0.5
        w_heading = 0.3
        w_speed = 0.2

        # 距离减少奖励
        if self.last_distance is not None:
            distance_reduction = max(0, self.last_distance - current_distance)
            reward += distance_reduction * 0.5

        # 路径偏差惩罚
        reward -= w_deviation * path_deviation * 0.5

        # 路径角度惩罚
        reward -= w_deviation * path_angle_error * 0.5

        # 高度奖励/惩罚
        alt_threshold = 5.0  # 严格高度误差阈值
        if abs(alt_error) < alt_threshold:
            reward += w_altitude * 200 * (1 - abs(alt_error) / alt_threshold)
        else:
            reward -= w_altitude * abs(alt_error) * 0.8

        # 航向奖励
        if heading_error < 5:
            reward += w_heading * 50 * (1 - heading_error / 5)
        else:
            reward -= w_heading * heading_error * 0.5

        # 速度奖励
        if speed_error < 10:
            reward += w_speed * 50 * (1 - speed_error / 10)
        else:
            reward -= w_speed * speed_error * 0.3

        # 攻角惩罚
        if abs(data["AoA"]) > AOA_THRESHOLD:
            reward -= 100

        # 速度范围惩罚
        if data["IAS"] < SPEED_LOWER_LIMIT or data["IAS"] > SPEED_UPPER_LIMIT:
            reward -= 50

        # 坠毁惩罚
        if data["LLA_balt"] < CRASH_ALT_THRESHOLD:
            reward -= 2000

        # 偏离过远惩罚
        if current_distance > INTERCEPT_DISTANCE * 1.5:
            reward -= 200

        # 时间惩罚
        reward -= 5.0

        return reward, data["LLA_balt"] < CRASH_ALT_THRESHOLD or current_distance > INTERCEPT_DISTANCE * 1.5

    def _is_done(self, state, current_distance, current_altitude):
        """判断是否终止"""
        if current_altitude < CRASH_ALT_THRESHOLD:
            logger.warning("Crash detected!")
            return True
        if current_distance > INTERCEPT_DISTANCE * 1.5:
            logger.warning("Too far from runway!")
            return True
        if abs(state[6]) > AOA_THRESHOLD:
            logger.warning("AoA threshold exceeded!")
            return True
        if state[3] < SPEED_LOWER_LIMIT or state[3] > SPEED_UPPER_LIMIT:
            logger.warning("Speed out of range!")
            return True
        if current_distance < 100 and abs(current_altitude - RUNWAY_ALT) < 50 and state[3] < LANDING_SPEED * 1.2:
            logger.info("Successful landing!")
            self._plot_trajectory()
            return True
        if self.steps >= 1000:
            logger.warning("Time limit reached!")
            return True
        return False

    def _take_action(self, action, data):
        """执行动作"""
        pitch_action, roll_action, throttle_action, flaps_action, airbrake_action, foot_action = action

        # 平滑动作
        for i in range(6):
            self.smoothed_action[i] = 0.7 * self.last_action[i] + 0.3 * action[i]
            self.last_action[i] = self.smoothed_action[i]
        pitch_action, roll_action, throttle_action, flaps_action, airbrake_action, foot_action = self.smoothed_action

        # 动态调整动作持续时间
        duration = 0.2 * max(0.5, min(1.5, data["IAS"] / APPROACH_SPEED)) * max(1.0, min(1.5, data["AoA"] / 10))

        # 俯仰控制
        if pitch_action > 0.1:
            pitch_down(duration)  # 提升机头
        elif pitch_action < -0.1:
            pitch_up(duration)  # 降低机头

        # 滚转控制
        if roll_action > 0.1:
            roll_right(duration)
        elif roll_action < -0.1:
            roll_left(duration)

        # 油门控制
        if throttle_action > 0.1:
            throttle_up(duration)
        elif throttle_action < -0.1:
            throttle_down(duration)

        # 襟翼控制
        if flaps_action > 0.1 and data["LLA_balt"] - RUNWAY_ALT < LOW_ALT_THRESHOLD and data["Flaps"] < 0.9:
            flapsland_control(0.1)

        # 减速板控制
        desired_airbrake = 1 if airbrake_action > 0.1 and data["IAS"] > SPEED_UPPER_LIMIT else 0
        if (desired_airbrake == 1 and data.get("AirBrake", 0) < 0.5) or (
                desired_airbrake == 0 and data.get("AirBrake", 0) > 0.5):
            airbrake_control(0.1)

        # 脚舵控制
        if foot_action > 0.1:
            foot_left(duration)
        elif foot_action < -0.1:
            foot_right(duration)

        # 起落架控制
        if data["LLA_balt"] - RUNWAY_ALT < LOW_ALT_THRESHOLD and data["Gears"] < 0.9:
            gearson_control(0.1)

        # 高攻角恢复
        if data["AoA"] > AOA_THRESHOLD * 0.8:
            throttle_up(0.2)
            pitch_up(0.2)
            logger.warning("High AoA detected, applying recovery actions")

    def step(self, action):
        """执行一步"""
        try:
            activate_dcs_window()
        except Exception as e:
            logger.error(f"Failed to activate DCS window in step: {str(e)}")

        data = self._get_dcs_data_with_retry()
        if data is None:
            data = self._get_fallback_data()

        self._take_action(action, data)

        # 动态控制频率
        sleep_time = 0.1 if data["IAS"] < 60 else 0.05
        time.sleep(sleep_time)

        data = self._get_dcs_data_with_retry()
        if data is None:
            data = self._get_fallback_data()

        new_state = self._extract_state(data)
        current_distance = geodesic(
            (data["LLA_lat"], data["LLA_long"]), self.runway_point
        ).meters
        current_altitude = data["LLA_balt"]

        reward, done_flag = self._calculate_reward(new_state, current_distance, data)
        done = self._is_done(new_state, current_distance, current_altitude) or done_flag

        self.total_reward += reward
        self.state = new_state
        self.last_distance = current_distance
        self.last_altitude = current_altitude
        self.steps += 1

        self.flight_trajectory.append((data["LLA_lat"], data["LLA_long"], data["LLA_balt"]))
        self._write_flight_data(data)
        self._write_deviation_data(data)

        logger.info(
            f"Episode: {self.current_episode}, Step: {self.steps}, "
            f"Alt: {current_altitude:.1f}m, IAS: {data['IAS']:.1f}m/s, "
            f"AoA: {data['AoA']:.1f}°, VS: {data['VV'][2]:.2f}m/s, "
            f"Reward: {reward:.2f}, Total Reward: {self.total_reward:.2f}"
        )

        # 定期或终止时绘制轨迹
        if done or self.steps % 50 == 0:
            self._plot_trajectory()

        # 记录训练日志
        with open(TRAINING_LOG_CSV, mode='a', newline='') as file:
            writer = csv.writer(file)
            writer.writerow([
                self.current_episode, self.steps, reward, self.total_reward,
                current_altitude, data["IAS"], new_state[8], data["AoA"],
                data["VV"][2], data["Gears"], data["Flaps"]
            ])

        return new_state, reward, done, {}

    def _write_flight_data(self, data):
        """记录飞行数据到CSV"""
        dist_to_runway = geodesic(
            (data["LLA_lat"], data["LLA_long"]), self.runway_point
        ).meters
        with open(FLIGHT_DATA_CSV, mode='a', newline='') as file:
            writer = csv.writer(file)
            writer.writerow([
                self.current_episode, self.steps,
                data["LLA_long"], data["LLA_lat"], data["LLA_balt"],
                data["IAS"], data["AoA"], data["pitch_deg"], data["bank_deg"],
                data["heading_deg"], data["VV"][2], data["Gears"], data["Flaps"],
                data["AirBrake"], data["Fuel"], dist_to_runway
            ])

    def _write_deviation_data(self, data):
        """记录飞机与下滑道偏差到CSV"""
        current_pos = (data["LLA_lat"], data["LLA_long"])
        nearest_point, path_deviation = self._find_nearest_glide_point(current_pos)
        alt_error = data["LLA_balt"] - nearest_point["alt"]

        target_vec = np.array([nearest_point["lat"] - current_pos[0], nearest_point["lon"] - current_pos[1]])
        heading_vec = np.array(
            [math.cos(math.radians(data["heading_deg"])), math.sin(math.radians(data["heading_deg"]))])
        path_angle_error = math.degrees(math.acos(np.clip(
            np.dot(target_vec, heading_vec) / (np.linalg.norm(target_vec) * np.linalg.norm(heading_vec) + 1e-6), -1, 1
        )))

        with open(GLIDE_DEVIATION_CSV, mode='a', newline='') as file:
            writer = csv.writer(file)
            writer.writerow([
                self.current_episode, self.steps,
                data["LLA_long"], data["LLA_lat"], data["LLA_balt"],
                path_deviation, alt_error, path_angle_error
            ])